import os
from adalflow.core import Component, Generator, Prompt
from adalflow.components.model_client import BedrockAPIClient

PROFILE_NAME = os.getenv("PROFILE_NAME")
REGION_NAME = os.getenv("REGION_NAME",)

class SimpleQA(Component):
    def __init__(self):
        super().__init__()
        template = r"""<SYS>
        You are a helpful assistant.
        </SYS>
        User: {{input_str}}
        You:
        """
        self.generator = Generator(
            model_client=BedrockAPIClient(
                aws_profile_name=PROFILE_NAME, 
                aws_region_name=REGION_NAME), 
            model_kwargs={
                "modelId": "anthropic.claude-3-sonnet-20240229-v1:0",
                "inferenceConfig": {
                    "temperature": 0.8
                }
            }, template=template
        )

    def call(self, query):
        return self.generator({"input_str": query})


qa = SimpleQA()
answer = qa("What is LightRAG?")
print(answer)